

      
#!/usr/bin/env python3
"""
简单的Octo PyTorch推理示例
使用test_conversion目录中的权重进行快速推理测试
"""

import torch
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import requests

from octo.model.octo_model_pt import OctoModelPt
import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
os.environ['HF_TOKEN'] = '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'


def main():
    print("加载Octo PyTorch模型...")
    
    loaded_dict = OctoModelPt.load_pretrained("test_conversion")
    model = loaded_dict['octo_model']
    
    # device = "cuda" if torch.cuda.is_available() else "cpu"
    device = 'cuda'
    print(f'use {device}')
    model = model.to(device)
    model.eval()
    
    print(f"模型已加载到设备: {device}")
    
    print("下载测试图像...")
    IMAGE_URL = "https://rail.eecs.berkeley.edu/datasets/bridge_release/raw/bridge_data_v2/datacol2_toykitchen7/drawer_pnp/01/2023-04-19_09-18-15/raw/traj_group0/traj0/images0/im_12.jpg"
    
    try:
        img = np.array(Image.open(requests.get(IMAGE_URL, stream=True).raw).resize((256, 256)))
        print("图像下载成功")
    except:
        print("使用随机图像")
        img = np.random.randint(0, 256, (256, 256, 3), dtype=np.uint8)
    
    plt.figure(figsize=(6, 6))
    plt.imshow(img)
    plt.title("输入图像")
    plt.axis('off')
    plt.savefig('input.png')
    plt.show()
    
    print("创建任务和观察...")
    tasks = model.create_tasks(texts=["pick up the fork"], device=device)
    
    img_tensor = torch.from_numpy(img.transpose((2, 0, 1))).unsqueeze(0).unsqueeze(0).to(device)
    observation = {
        "image_primary": img_tensor,
        "timestep_pad_mask": torch.tensor([[True]], device=device)
    }
    
    print("执行推理...")
    with torch.no_grad():
        actions = model.sample_actions(
            observation,
            tasks,
            unnormalization_statistics=model.dataset_statistics["bridge_dataset"]["action"],
            generator=torch.Generator(device).manual_seed(42)
        )
    
    print(f"预测动作形状: {actions.shape}")
    print(f"预测动作值:")
    action_labels = ['x', 'y', 'z', 'yaw', 'pitch', 'roll', 'grasp']
    action_values = actions[0, 0].cpu().numpy()
    
    for label, value in zip(action_labels, action_values):
        print(f"  {label}: {value:.4f}")
    
    plt.figure(figsize=(10, 6))
    bars = plt.bar(action_labels, action_values)
    plt.title("预测的机器人动作")
    plt.ylabel("动作值")
    plt.grid(True, alpha=0.3)
    
    for bar, val in zip(bars, action_values):
        plt.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01,
                f'{val:.3f}', ha='center', va='bottom')
    
    plt.tight_layout()
    plt.savefig('actions.png')
    plt.show()
    
    print("推理完成! 结果已保存为 input.png 和 actions.png")


if __name__ == "__main__":
    main()

    


